import numpy as np
from utils import im2col, col2im


class Convolution:
    def __init__(self, input_dim=(1, 28, 28), num=30, size=5, stride=1, pad=0, weight_init_std=0.01):
        self.W = weight_init_std * \
                 np.random.randn(num, input_dim[0], size, size)
        self.b = np.zeros(num)
        self.stride = stride
        self.pad = pad

        # 中间数据（backward时使用）
        self.x = None
        self.col = None
        self.col_W = None

        # 权重和偏置参数的梯度
        self.dW = None
        self.db = None

    def forward(self, x):
        FN, C, FH, FW = self.W.shape  # 卷积核的形状
        N, C, H, W = x.shape  # 输入数据形状
        ###请补充代码获取输出数据的高和宽分别存储在变量out_h和变量out_w
        out_h = (H + 2 * self.pad - FH) // self.stride + 1
        out_w = (W + 2 * self.pad - FW) // self.stride + 1
        ###请补充代码来完成卷积运算
        ###(其中col、col_W、out分别表示展开的数据、卷积核、卷积运算的结果)
        col = im2col(x, FH, FW, self.stride, self.pad)
        col_W = self.W.reshape(-1, FN)
        out = np.dot(col, col_W)

        # 输出大小转换为合适的形状
        # transpose会更改多维数组的轴的顺序，将输出数据形状由(N,H,W,C)转变为(N,C,H,W)
        out = out.reshape(N, out_h, out_w, -1).transpose(0, 3, 1, 2)

        # 更新backward过程需要用到的中间数据
        self.x = x
        self.col = col
        self.col_W = col_W

        return out

    def backward(self, dout):
        FN, C, FH, FW = self.W.shape
        dout = dout.transpose(0, 2, 3, 1).reshape(-1, FN)

        self.db = np.sum(dout, axis=0)
        self.dW = np.dot(self.col.T, dout)
        self.dW = self.dW.transpose(1, 0).reshape(FN, C, FH, FW)

        dcol = np.dot(dout, self.col_W.T)
        dx = col2im(dcol, self.x.shape, FH, FW, self.stride, self.pad)

        return dx
